#!/usr/bin/env python
# -*- coding: utf-8 -*-
# @Author: slp
# @Date:   2018-05-04 15:10:05
# @Last Modified by:   ystlong
# @Last Modified time: 2018-11-25 18:52:11

import urllib
import urllib2
import sys
import re
import os
import json
import time
import uuid

def http_conn(url, data, headers={}):
	request = urllib2.Request(url, data)
	for k in headers:
		request.add_header(k, headers[k])
	request.add_header("Content-Type", "application/json")
	response = urllib2.urlopen(request)
	if response.getcode() < 200 or response.getcode() >= 300:
		raise Exception("http request not success, code:%d"%(response.getcode()))
	return response

def read_mesos_data(response, _json=True):
	length = []
	while True:
		r = response.read(1)
		if r != '\n':
			length.append(r)
		else:
			break
	length = int("".join(length))
	data = response.read(length)
	if _json:
		return json.loads(data)
	return data

def subscribe(url):
	data = {}
	data["type"] = "SUBSCRIBE"
	# data["subscribe"] = {"framework_info":{"user":"foo","name":"example http", "roles":["test"], "capabilities":[{"type":"MULTI_ROLE"}]}}
	data["subscribe"] = {"framework_info":{"user":"slp","name":"example http", "roles":[], "capabilities":[{"type":"RESERVATION_REFINEMENT"}]}}
	data = json.dumps(data)
	response = http_conn(url, data, headers={"Accept":"application/json"})
	# print data
	stream_id = response.headers["Mesos-Stream-Id"]
	subscribe = None
	offers = None
	offerd = False
	while True:
		data = read_mesos_data(response)
		print "****************", data
		# print data
		if type(data) != dict:
			print "invaild response"
			break
		elif data["type"] == "SUBSCRIBED":
			subscribe = data["subscribed"]
		elif data["type"] == "OFFERS":
			offers = data["offers"]["offers"]
			# for offer in offers:
			# 	print offer.keys()
			# 	print offer
			if not offerd:
				accept(url, subscribe, offers, stream_id)
				offerd = True
		elif data["type"] == "HEARTBEAT":
			print data
			pass
		else:
			print "=========", data

def new_operation():
	operation_tpl = {
		"type" : "LAUNCH",
		"launch" : {
			"task_infos" : []
		}
	}
	return operation_tpl			

def new_task(name, agent_id, framework_id):
	task_tpl = {
		"name"   : "My task",
		"task_id" : {"value" : "4c35a7fd-78d9-4eb0-9da0-9eecb1c965c3-my-task"},
		"agent_id" : {"value": ""},
		"executor" : {
			"command" : {
				"shell" : True,
				# "value": "whoami; sleep 120"
				"value": "~/test.sh"
			},
			"executor_id" : {"value": "4c35a7fd-78d9-4eb0-9da0-9eecb1c965c3-my-executor"},
		},
		"resources" : [
			{
				"allocation_info" : {"role": "*"},
				"name" : "cpus",
				"role" : "*",
				"type" : "SCALAR",
				"scalar" : {"value": 1.0}
			},
			{
				"allocation_info" : {"role": "*"},
				"name" : "mem",
				"role" : "*",
				"type" : "SCALAR",
				"scalar" : {"value": 128.0}
			}
		]
	}
	task_tpl["name"] = name
	task_tpl["agent_id"]["value"] = agent_id
	task_tpl["task_id"]["value"] = str(uuid.uuid4())
	task_tpl["executor"]["executor_id"]["value"] = str(uuid.uuid4())
	return task_tpl


def accept_curl(url, stream_id, data):
	curl_tpl = '''curl "$url" -H "Content-Type: application/json" -H "Mesos-Stream-Id: $mesos_header_id" --data @$json_file'''
	tmp_json_file = ".json.tmp"
	with open(tmp_json_file, "w") as f:
		f.write(data)
	k = "Mesos-Stream-Id"
	cmd = curl_tpl.replace("$url", url).replace("$mesos_header_id", stream_id).replace("$json_file", tmp_json_file)
	print cmd
	os.system(cmd)

def acknowledge(framework_id, agent_id, task_id, uuid):
	ack_tpl = {
		"framework_id": {"value": framework_id},
		"type": "ACKNOWLEDGE",
		"acknowledge": {
			"agent_id": agent_id,
			"task_id": task_id,
			"uuid": uuid
		}
	}
	return ack_tpl

def accept(url, subscribe, offers, stream_id):
	# accept_curl(url, framwork, stream_id)
	# return
	offer = offers[0]
	json_file = "accept.json"
	api_info = None
	with open(json_file) as f:
		api_info = json.load(f)

	framework_id = offer["framework_id"]["value"]
	agent_id = offer["agent_id"]["value"]

	# api_info["framework_id"]["value"] = subscribe["subscribed"]["framework_id"]["value"]
	api_info["framework_id"]["value"] = framework_id
	api_info["accept"]["offer_ids"] = [{"value":offer["id"]["value"]}]
	operation = new_operation()
	task1 = new_task("task1", agent_id, framework_id)
	operation["launch"]["task_infos"].append(task1)

	api_info["accept"]["operations"] = []
	api_info["accept"]["operations"].append(operation)
	with open("task.json", "w") as f:
		json.dump(api_info, f)
	# api_info["operations"]
	# data = json.dumps(api_info, indent=True)
	data = json.dumps(api_info)
	# print data
	# accept_curl(url, stream_id, data)

	k = "Mesos-Stream-Id"
	response = http_conn(url, data, headers={k:stream_id})
	# data = read_mesos_data(response)
	# print "---------", data

def main():
	# return
	# url = "http://202.197.154.243:5050/api/v1/scheduler"
	url = "http://172.17.0.2:5050/api/v1/scheduler"
	subscribe(url)
	os.system("sleep 3")

if __name__ == '__main__':
	main()
